import cn.wxw.ai.common.Constants;
import io.milvus.common.clientenum.FunctionType;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.*;

/**
 * @description: 创建集合 collection 与索引 index
 * @author: wxw
 * @date: 2025/9/8
 */
@Component
@Slf4j
public class CollectionManager {

    @Resource
    private MilvusClientV2 client;

    public void createCollection(String collectionType, String collectionName) throws Exception {
        try {
            // 创建密集向量 Collection
            if (collectionType == Constants.COLLECTION_TYPE.VECTOR.getCode()) {
                final CreateCollectionReq vectorCollectionReq = CreateCollectionReq.builder()
                        .collectionName(collectionName)
                        .collectionSchema(this.createVectorCollectionSchema())
                        .indexParams(this.createVectorIndex())
                        .build();
                client.createCollection(vectorCollectionReq);
                log.info("Successfully created vector collection: {}", collectionName);
            }

            // 创建稀疏向量 Collection
            if (collectionType == Constants.COLLECTION_TYPE.TEXT.getCode()) {
                final CreateCollectionReq textCollectionReq = CreateCollectionReq.builder()
                        .collectionName(collectionName)
                        .collectionSchema(this.createTextCollectionSchema())
                        .indexParams(this.createTextIndex())
                        .build();
                client.createCollection(textCollectionReq);
                log.info("Successfully created text collection: {}", collectionName);
            }

            // 创建多向量 Collection
            if (collectionType == Constants.COLLECTION_TYPE.MULTI_VECTOR.getCode()) {
                final CreateCollectionReq vectorCollectionReq = CreateCollectionReq.builder()
                        .collectionName(collectionName)
                        .collectionSchema(this.createMultiVectorCollectionSchema())
                        .indexParams(this.createMultiVectorIndex())
                        .build();
                client.createCollection(vectorCollectionReq);
                log.info("Successfully created multi vector collection: {}", collectionName);
            }
        } catch (Exception e) {
            log.error("Error creating collections: {}", e.getMessage(), e);
            throw new Exception("Failed to create collections", e);
        }
    }
    /**
     * 检查集合是否存在
     *
     * @param collectionName 集合名称
     * @return 是否存在
     */
    private boolean hasCollection(final String collectionName) {
        return client.hasCollection(HasCollectionReq.builder()
                .collectionName(collectionName)
                .build());
    }

    /**
     * 创建密集向量集合的 Schema
     * @return 向量集合 Schema
     */
    public CreateCollectionReq.CollectionSchema createVectorCollectionSchema() {
        final CreateCollectionReq.CollectionSchema schema = CreateCollectionReq.CollectionSchema.builder()
                .build();
        schema.addField(AddFieldReq.builder()
                .fieldName("id")
                .dataType(DataType.Int64)
                .isPrimaryKey(true)
                .autoID(true)
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("text")
                .dataType(DataType.VarChar)
                .maxLength(65535)
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("metadata")
                .dataType(DataType.JSON)
                .maxLength(65535)
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("text_dense")
                .dataType(DataType.FloatVector)
                .dimension(1536)
                .build());
        return schema;
    }

    /**
     * 创建稀疏向量集合的 Schema
     * id主键自增
     * 配置了中文分词器和 BM25 函数
     * @return 文本集合 Schema
     */
    public CreateCollectionReq.CollectionSchema createTextCollectionSchema() {
        final CreateCollectionReq.CollectionSchema schema = CreateCollectionReq.CollectionSchema.builder()
                .build();
        schema.addField(AddFieldReq.builder()
                .fieldName("id")
                .dataType(DataType.Int64)
                .isPrimaryKey(true)
                .autoID(true)
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("text")
                .dataType(DataType.VarChar)
                .maxLength(65535)
                .enableAnalyzer(true)
                .analyzerParams(Map.of("type", "chinese")) // 中文分词器
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("text_sparse")
                .dataType(DataType.SparseFloatVector)
                .build());

        schema.addFunction(CreateCollectionReq.Function.builder()
                .functionType(FunctionType.BM25)
                .name("text_bm25_emb")
                .inputFieldNames(Collections.singletonList("text"))
                .outputFieldNames(Collections.singletonList("text_sparse"))
                .build());
        return schema;
    }


    /**
     * 创建多向量(存储文本的密集向量/存储文本的稀疏向量) Collection
     * @return 向量集合 Schema
     */
    public CreateCollectionReq.CollectionSchema createMultiVectorCollectionSchema() {
        final CreateCollectionReq.CollectionSchema schema = CreateCollectionReq.CollectionSchema.builder()
                .build();
        schema.addField(AddFieldReq.builder()
                .fieldName("id")
                .dataType(DataType.Int64)
                .isPrimaryKey(true)
                .autoID(true)
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("text")
                .dataType(DataType.VarChar)
                .maxLength(65535)
                .enableAnalyzer(true)
                .analyzerParams(Map.of("type", "chinese")) // 中文分词器
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("text_dense")
                .dataType(DataType.FloatVector)
                .dimension(1536)
                .build());
        schema.addField(AddFieldReq.builder()
                .fieldName("text_sparse")
                .dataType(DataType.SparseFloatVector)
                .build());
        schema.addFunction(CreateCollectionReq.Function.builder()
                .functionType(FunctionType.BM25)
                .name("text_bm25_emb")
                .inputFieldNames(Collections.singletonList("text"))
                .outputFieldNames(Collections.singletonList("text_sparse"))
                .build());
        return schema;
    }

    /**
     * 创建密集向量检索的索引参数
     * 使用 HNSW 索引类型，余弦相似度度量方式
     * 配置了 M=8 和 efConstruction=64 的 HNSW 参数
     *
     * @return 向量索引参数列表
     */
    public List<IndexParam> createVectorIndex() {
        final List<IndexParam> indexes = new ArrayList<>();
        final Map<String, Object> extraParams = new HashMap<>();
        extraParams.put("M", 8);
        extraParams.put("efConstruction", 64);
        indexes.add(IndexParam.builder()
                .fieldName("text_dense")
                .indexType(IndexParam.IndexType.HNSW)
                .metricType(IndexParam.MetricType.COSINE)
                .extraParams(extraParams)
                .build());
        return indexes;
    }

    /**
     * 创建稀疏向量检索的索引参数
     * 使用 AUTOINDEX 索引类型，BM25 相似度度量方式
     *
     * @return 文本索引参数列表
     */
    public List<IndexParam> createTextIndex() {
        final List<IndexParam> indexes = new ArrayList<>();
        indexes.add(IndexParam.builder()
                .fieldName("text_sparse")
                .indexType(IndexParam.IndexType.AUTOINDEX)
                .metricType(IndexParam.MetricType.BM25)
                .build());
        return indexes;
    }

    /**
     * 创建多向量检索的索引参数
     *
     * @return 索引参数列表
     */
    public List<IndexParam> createMultiVectorIndex() {

        IndexParam indexParamForTextDense = IndexParam.builder()
                .fieldName("text_dense")
                .indexName("text_dense_index")
                .indexType(IndexParam.IndexType.AUTOINDEX)
                .metricType(IndexParam.MetricType.IP)
                .build();

        Map<String, Object> sparseParams = new HashMap<>();
        sparseParams.put("inverted_index_algo", "DAAT_MAXSCORE");
        IndexParam indexParamForTextSparse = IndexParam.builder()
                .fieldName("text_sparse")
                .indexName("text_sparse_index")
                .indexType(IndexParam.IndexType.SPARSE_INVERTED_INDEX)
                .metricType(IndexParam.MetricType.BM25)
                .extraParams(sparseParams)
                .build();

        List<IndexParam> indexParams = new ArrayList<>();
        indexParams.add(indexParamForTextDense);
        indexParams.add(indexParamForTextSparse);

        return indexParams;
    }

}
